{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial5: Aggregation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this tutorial we will override the aggregation method of the GIN convolution module of Pytorch Geometric implementing the following methods:\n",
    "\n",
    "- Principal Neighborhood Aggregation (PNA)\n",
    "- Learning Aggregation Functions (LAF)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11.0\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "os.environ['TORCH'] = torch.__version__\n",
    "print(torch.__version__)\n",
    "\n",
    "!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
    "!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
    "!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f4eec098e70>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Message Passing Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import MessagePassing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['T_destination',\n",
       " '__annotations__',\n",
       " '__call__',\n",
       " '__check_input__',\n",
       " '__class__',\n",
       " '__collect__',\n",
       " '__delattr__',\n",
       " '__dict__',\n",
       " '__dir__',\n",
       " '__doc__',\n",
       " '__eq__',\n",
       " '__format__',\n",
       " '__ge__',\n",
       " '__getattr__',\n",
       " '__getattribute__',\n",
       " '__gt__',\n",
       " '__hash__',\n",
       " '__init__',\n",
       " '__init_subclass__',\n",
       " '__le__',\n",
       " '__lift__',\n",
       " '__lt__',\n",
       " '__module__',\n",
       " '__ne__',\n",
       " '__new__',\n",
       " '__reduce__',\n",
       " '__reduce_ex__',\n",
       " '__repr__',\n",
       " '__set_size__',\n",
       " '__setattr__',\n",
       " '__setstate__',\n",
       " '__sizeof__',\n",
       " '__str__',\n",
       " '__subclasshook__',\n",
       " '__weakref__',\n",
       " '_apply',\n",
       " '_call_impl',\n",
       " '_get_backward_hooks',\n",
       " '_get_name',\n",
       " '_load_from_state_dict',\n",
       " '_maybe_warn_non_full_backward_hook',\n",
       " '_named_members',\n",
       " '_register_load_state_dict_pre_hook',\n",
       " '_register_state_dict_hook',\n",
       " '_replicate_for_data_parallel',\n",
       " '_save_to_state_dict',\n",
       " '_slow_forward',\n",
       " '_version',\n",
       " 'add_module',\n",
       " 'aggregate',\n",
       " 'apply',\n",
       " 'bfloat16',\n",
       " 'buffers',\n",
       " 'children',\n",
       " 'combine',\n",
       " 'cpu',\n",
       " 'cuda',\n",
       " 'double',\n",
       " 'dump_patches',\n",
       " 'edge_update',\n",
       " 'edge_updater',\n",
       " 'eval',\n",
       " 'explain',\n",
       " 'explain_message',\n",
       " 'extra_repr',\n",
       " 'float',\n",
       " 'forward',\n",
       " 'get_buffer',\n",
       " 'get_extra_state',\n",
       " 'get_parameter',\n",
       " 'get_submodule',\n",
       " 'half',\n",
       " 'jittable',\n",
       " 'load_state_dict',\n",
       " 'message',\n",
       " 'message_and_aggregate',\n",
       " 'modules',\n",
       " 'named_buffers',\n",
       " 'named_children',\n",
       " 'named_modules',\n",
       " 'named_parameters',\n",
       " 'parameters',\n",
       " 'propagate',\n",
       " 'register_aggregate_forward_hook',\n",
       " 'register_aggregate_forward_pre_hook',\n",
       " 'register_backward_hook',\n",
       " 'register_buffer',\n",
       " 'register_edge_update_forward_hook',\n",
       " 'register_edge_update_forward_pre_hook',\n",
       " 'register_forward_hook',\n",
       " 'register_forward_pre_hook',\n",
       " 'register_full_backward_hook',\n",
       " 'register_message_and_aggregate_forward_hook',\n",
       " 'register_message_and_aggregate_forward_pre_hook',\n",
       " 'register_message_forward_hook',\n",
       " 'register_message_forward_pre_hook',\n",
       " 'register_module',\n",
       " 'register_parameter',\n",
       " 'register_propagate_forward_hook',\n",
       " 'register_propagate_forward_pre_hook',\n",
       " 'requires_grad_',\n",
       " 'set_extra_state',\n",
       " 'share_memory',\n",
       " 'special_args',\n",
       " 'state_dict',\n",
       " 'to',\n",
       " 'to_empty',\n",
       " 'train',\n",
       " 'type',\n",
       " 'update',\n",
       " 'xpu',\n",
       " 'zero_grad']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dir(MessagePassing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are interested in the <span style='color:Blue'>aggregate</span> method, or, if you are using a sparse adjacency matrix, in the <span style='color:Blue'>message_and_aggregate</span> method. Convolutional classes in PyG extend MessagePassing, we construct our custom convoutional class extending GINConv."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Scatter operation in <span style='color:Blue'>aggregate</span>:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/add.svg?sanitize=true\" width=\"500\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.nn import Parameter, Module, Sigmoid\n",
    "import torch\n",
    "import torch_scatter\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class AbstractLAFLayer(Module):\n",
    "    def __init__(self, **kwargs):\n",
    "        super(AbstractLAFLayer, self).__init__()\n",
    "        assert 'units' in kwargs or 'weights' in kwargs\n",
    "        if 'device' in kwargs.keys():\n",
    "            self.device = kwargs['device']\n",
    "        else:\n",
    "            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "        self.ngpus = torch.cuda.device_count()\n",
    "        \n",
    "        if 'kernel_initializer' in kwargs.keys():\n",
    "            assert kwargs['kernel_initializer'] in [\n",
    "                'random_normal',\n",
    "                'glorot_normal',\n",
    "                'he_normal',\n",
    "                'random_uniform',\n",
    "                'glorot_uniform',\n",
    "                'he_uniform']\n",
    "            self.kernel_initializer = kwargs['kernel_initializer']\n",
    "        else:\n",
    "            self.kernel_initializer = 'random_normal'\n",
    "\n",
    "        if 'weights' in kwargs.keys():\n",
    "            self.weights = Parameter(kwargs['weights'].to(self.device), \\\n",
    "                                     requires_grad=True)\n",
    "            self.units = self.weights.shape[1]\n",
    "        else:\n",
    "            self.units = kwargs['units']\n",
    "            params = torch.empty(12, self.units, device=self.device)\n",
    "            if self.kernel_initializer == 'random_normal':\n",
    "                torch.nn.init.normal_(params)\n",
    "            elif self.kernel_initializer == 'glorot_normal':\n",
    "                torch.nn.init.xavier_normal_(params)\n",
    "            elif self.kernel_initializer == 'he_normal':\n",
    "                torch.nn.init.kaiming_normal_(params)\n",
    "            elif self.kernel_initializer == 'random_uniform':\n",
    "                torch.nn.init.uniform_(params)\n",
    "            elif self.kernel_initializer == 'glorot_uniform':\n",
    "                torch.nn.init.xavier_uniform_(params)\n",
    "            elif self.kernel_initializer == 'he_uniform':\n",
    "                torch.nn.init.kaiming_uniform_(params)\n",
    "            self.weights = Parameter(params, \\\n",
    "                                     requires_grad=True)\n",
    "        e = torch.tensor([1,-1,1,-1], dtype=torch.float32, device=self.device)\n",
    "        self.e = Parameter(e, requires_grad=False)\n",
    "        num_idx = torch.tensor([1,1,0,0], dtype=torch.float32, device=self.device).\\\n",
    "                                view(1,1,-1,1)\n",
    "        self.num_idx = Parameter(num_idx, requires_grad=False)\n",
    "        den_idx = torch.tensor([0,0,1,1], dtype=torch.float32, device=self.device).\\\n",
    "                                view(1,1,-1,1)\n",
    "        self.den_idx = Parameter(den_idx, requires_grad=False)\n",
    "        \n",
    "\n",
    "class LAFLayer(AbstractLAFLayer):\n",
    "    def __init__(self, eps=1e-7, **kwargs):\n",
    "        super(LAFLayer, self).__init__(**kwargs)\n",
    "        self.eps = eps\n",
    "    \n",
    "    def forward(self, data, index, dim=0, **kwargs):\n",
    "        eps = self.eps\n",
    "        sup = 1.0 - eps \n",
    "        e = self.e\n",
    "\n",
    "        x = torch.clamp(data, eps, sup)\n",
    "        x = torch.unsqueeze(x, -1)\n",
    "        e = e.view(1,1,-1)        \n",
    "\n",
    "        exps = (1. - e)/2. + x*e \n",
    "        exps = torch.unsqueeze(exps, -1)\n",
    "        exps = torch.pow(exps, torch.relu(self.weights[0:4]))\n",
    "\n",
    "        scatter = torch_scatter.scatter_add(exps, index.view(-1), dim=dim)\n",
    "        scatter = torch.clamp(scatter, eps)\n",
    "\n",
    "        sqrt = torch.pow(scatter, torch.relu(self.weights[4:8]))\n",
    "        alpha_beta = self.weights[8:12].view(1,1,4,-1)\n",
    "        terms = sqrt * alpha_beta\n",
    "\n",
    "        num = torch.sum(terms * self.num_idx, dim=2)\n",
    "        den = torch.sum(terms * self.den_idx, dim=2)\n",
    "        \n",
    "        multiplier = 2.0*torch.clamp(torch.sign(den), min=0.0) - 1.0\n",
    "\n",
    "        den = torch.where((den < eps) & (den > -eps), multiplier*eps, den)\n",
    "\n",
    "        res = num / den\n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import GINConv\n",
    "from torch.nn import Linear"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LAF Aggregation Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"laf.png\" width=\"500\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GINLAFConv(GINConv):\n",
    "    def __init__(self, nn, units=1, node_dim=32, **kwargs):\n",
    "        super(GINLAFConv, self).__init__(nn, **kwargs)\n",
    "        self.laf = LAFLayer(units=units, kernel_initializer='random_uniform')\n",
    "        self.mlp = torch.nn.Linear(node_dim*units, node_dim)\n",
    "        self.dim = node_dim\n",
    "        self.units = units\n",
    "    \n",
    "    def aggregate(self, inputs, index):\n",
    "        x = torch.sigmoid(inputs)\n",
    "        x = self.laf(x, index)\n",
    "        x = x.view((-1, self.dim * self.units))\n",
    "        x = self.mlp(x)\n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PNA Aggregation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"pna.png\" width=\"800\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GINPNAConv(GINConv):\n",
    "    def __init__(self, nn, node_dim=32, **kwargs):\n",
    "        super(GINPNAConv, self).__init__(nn, **kwargs)\n",
    "        self.mlp = torch.nn.Linear(node_dim*12, node_dim)\n",
    "        self.delta = 2.5749\n",
    "    \n",
    "    def aggregate(self, inputs, index):\n",
    "        sums = torch_scatter.scatter_add(inputs, index, dim=0)\n",
    "        maxs = torch_scatter.scatter_max(inputs, index, dim=0)[0]\n",
    "        means = torch_scatter.scatter_mean(inputs, index, dim=0)\n",
    "        var = torch.relu(torch_scatter.scatter_mean(inputs ** 2, index, dim=0) - means ** 2)\n",
    "        \n",
    "        aggrs = [sums, maxs, means, var]\n",
    "        c_idx = index.bincount().float().view(-1, 1)\n",
    "        l_idx = torch.log(c_idx + 1.)\n",
    "        \n",
    "        amplification_scaler = [c_idx / self.delta * a for a in aggrs]\n",
    "        attenuation_scaler = [self.delta / c_idx * a for a in aggrs]\n",
    "        combinations = torch.cat(aggrs+ amplification_scaler+ attenuation_scaler, dim=1)\n",
    "        x = self.mlp(combinations)\n",
    "    \n",
    "        return x\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the new classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_geometric.nn import MessagePassing, SAGEConv, GINConv, global_add_pool\n",
    "import torch_scatter\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import Sequential, Linear, ReLU\n",
    "from torch_geometric.datasets import TUDataset\n",
    "from torch_geometric.data import DataLoader\n",
    "import os.path as osp\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/antonio/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch_geometric/deprecation.py:12: UserWarning: 'data.DataLoader' is deprecated, use 'loader.DataLoader' instead\n",
      "  warnings.warn(out)\n"
     ]
    }
   ],
   "source": [
    "path = osp.join('./', 'data', 'TU')\n",
    "dataset = TUDataset(path, name='MUTAG').shuffle()\n",
    "test_dataset = dataset[:len(dataset) // 10]\n",
    "train_dataset = dataset[len(dataset) // 10:]\n",
    "test_loader = DataLoader(test_dataset, batch_size=128)\n",
    "train_loader = DataLoader(train_dataset, batch_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LAFNet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LAFNet, self).__init__()\n",
    "\n",
    "        num_features = dataset.num_features\n",
    "        dim = 32\n",
    "        units = 3\n",
    "        \n",
    "        nn1 = Sequential(Linear(num_features, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv1 = GINLAFConv(nn1, units=units, node_dim=num_features)\n",
    "        self.bn1 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn2 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv2 = GINLAFConv(nn2, units=units, node_dim=dim)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn3 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv3 = GINLAFConv(nn3, units=units, node_dim=dim)\n",
    "        self.bn3 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn4 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv4 = GINLAFConv(nn4, units=units, node_dim=dim)\n",
    "        self.bn4 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn5 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv5 = GINLAFConv(nn5, units=units, node_dim=dim)\n",
    "        self.bn5 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        self.fc1 = Linear(dim, dim)\n",
    "        self.fc2 = Linear(dim, dataset.num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "        x = self.bn2(x)\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "        x = self.bn3(x)\n",
    "        x = F.relu(self.conv4(x, edge_index))\n",
    "        x = self.bn4(x)\n",
    "        x = F.relu(self.conv5(x, edge_index))\n",
    "        x = self.bn5(x)\n",
    "        x = global_add_pool(x, batch)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PNANet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PNANet, self).__init__()\n",
    "\n",
    "        num_features = dataset.num_features\n",
    "        dim = 32\n",
    "\n",
    "        nn1 = Sequential(Linear(num_features, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv1 = GINPNAConv(nn1, node_dim=num_features)\n",
    "        self.bn1 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn2 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv2 = GINPNAConv(nn2, node_dim=dim)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn3 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv3 = GINPNAConv(nn3, node_dim=dim)\n",
    "        self.bn3 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn4 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv4 = GINPNAConv(nn4, node_dim=dim)\n",
    "        self.bn4 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn5 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv5 = GINPNAConv(nn5, node_dim=dim)\n",
    "        self.bn5 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        self.fc1 = Linear(dim, dim)\n",
    "        self.fc2 = Linear(dim, dataset.num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "        x = self.bn2(x)\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "        x = self.bn3(x)\n",
    "        x = F.relu(self.conv4(x, edge_index))\n",
    "        x = self.bn4(x)\n",
    "        x = F.relu(self.conv5(x, edge_index))\n",
    "        x = self.bn5(x)\n",
    "        x = global_add_pool(x, batch)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GINNet(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(GINNet, self).__init__()\n",
    "\n",
    "        num_features = dataset.num_features\n",
    "        dim = 32\n",
    "\n",
    "        nn1 = Sequential(Linear(num_features, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv1 = GINConv(nn1)\n",
    "        self.bn1 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn2 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv2 = GINConv(nn2)\n",
    "        self.bn2 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn3 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv3 = GINConv(nn3)\n",
    "        self.bn3 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn4 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv4 = GINConv(nn4)\n",
    "        self.bn4 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        nn5 = Sequential(Linear(dim, dim), ReLU(), Linear(dim, dim))\n",
    "        self.conv5 = GINConv(nn5)\n",
    "        self.bn5 = torch.nn.BatchNorm1d(dim)\n",
    "\n",
    "        self.fc1 = Linear(dim, dim)\n",
    "        self.fc2 = Linear(dim, dataset.num_classes)\n",
    "\n",
    "    def forward(self, x, edge_index, batch):\n",
    "        x = F.relu(self.conv1(x, edge_index))\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(self.conv2(x, edge_index))\n",
    "        x = self.bn2(x)\n",
    "        x = F.relu(self.conv3(x, edge_index))\n",
    "        x = self.bn3(x)\n",
    "        x = F.relu(self.conv4(x, edge_index))\n",
    "        x = self.bn4(x)\n",
    "        x = F.relu(self.conv5(x, edge_index))\n",
    "        x = self.bn5(x)\n",
    "        x = global_add_pool(x, batch)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.dropout(x, p=0.5, training=self.training)\n",
    "        x = self.fc2(x)\n",
    "        return F.log_softmax(x, dim=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 001, Train Loss: 0.6833775, Train Acc: 0.3294118, Test Acc: 0.3888889\n",
      "Epoch: 002, Train Loss: 0.7346479, Train Acc: 0.3294118, Test Acc: 0.3888889\n",
      "Epoch: 003, Train Loss: 0.5864959, Train Acc: 0.6705882, Test Acc: 0.6111111\n",
      "Epoch: 004, Train Loss: 0.6328506, Train Acc: 0.6705882, Test Acc: 0.6111111\n",
      "Epoch: 005, Train Loss: 0.5727439, Train Acc: 0.6705882, Test Acc: 0.6111111\n",
      "Epoch: 006, Train Loss: 0.5277878, Train Acc: 0.6705882, Test Acc: 0.6111111\n",
      "Epoch: 007, Train Loss: 0.4450684, Train Acc: 0.8294118, Test Acc: 0.8333333\n",
      "Epoch: 008, Train Loss: 0.5025334, Train Acc: 0.5058824, Test Acc: 0.4444444\n",
      "Epoch: 009, Train Loss: 0.5093456, Train Acc: 0.7941176, Test Acc: 0.8333333\n",
      "Epoch: 010, Train Loss: 0.4564617, Train Acc: 0.7411765, Test Acc: 0.9444444\n",
      "Epoch: 011, Train Loss: 0.5734278, Train Acc: 0.6823529, Test Acc: 0.6111111\n",
      "Epoch: 012, Train Loss: 0.4488182, Train Acc: 0.6764706, Test Acc: 0.6111111\n",
      "Epoch: 013, Train Loss: 0.4316032, Train Acc: 0.6764706, Test Acc: 0.7777778\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Input \u001b[0;32mIn [16]\u001b[0m, in \u001b[0;36m<cell line: 44>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     41\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m correct \u001b[38;5;241m/\u001b[39m \u001b[38;5;28mlen\u001b[39m(loader\u001b[38;5;241m.\u001b[39mdataset)\n\u001b[1;32m     44\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m101\u001b[39m):\n\u001b[0;32m---> 45\u001b[0m     train_loss \u001b[38;5;241m=\u001b[39m \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     46\u001b[0m     train_acc \u001b[38;5;241m=\u001b[39m test(train_loader)\n\u001b[1;32m     47\u001b[0m     test_acc \u001b[38;5;241m=\u001b[39m test(test_loader)\n",
      "Input \u001b[0;32mIn [16]\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(epoch)\u001b[0m\n\u001b[1;32m     22\u001b[0m data \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     23\u001b[0m optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m---> 24\u001b[0m output \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43medge_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdata\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     25\u001b[0m loss \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mnll_loss(output, data\u001b[38;5;241m.\u001b[39my)\n\u001b[1;32m     26\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36mLAFNet.forward\u001b[0;34m(self, x, edge_index, batch)\u001b[0m\n\u001b[1;32m     39\u001b[0m x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv4(x, edge_index))\n\u001b[1;32m     40\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn4(x)\n\u001b[0;32m---> 41\u001b[0m x \u001b[38;5;241m=\u001b[39m F\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconv5\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     42\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mbn5(x)\n\u001b[1;32m     43\u001b[0m x \u001b[38;5;241m=\u001b[39m global_add_pool(x, batch)\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch_geometric/nn/conv/gin_conv.py:74\u001b[0m, in \u001b[0;36mGINConv.forward\u001b[0;34m(self, x, edge_index, size)\u001b[0m\n\u001b[1;32m     71\u001b[0m     x: OptPairTensor \u001b[38;5;241m=\u001b[39m (x, x)\n\u001b[1;32m     73\u001b[0m \u001b[38;5;66;03m# propagate_type: (x: OptPairTensor)\u001b[39;00m\n\u001b[0;32m---> 74\u001b[0m out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpropagate\u001b[49m\u001b[43m(\u001b[49m\u001b[43medge_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msize\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msize\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     76\u001b[0m x_r \u001b[38;5;241m=\u001b[39m x[\u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m     77\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m x_r \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py:352\u001b[0m, in \u001b[0;36mMessagePassing.propagate\u001b[0;34m(self, edge_index, size, **kwargs)\u001b[0m\n\u001b[1;32m    349\u001b[0m         aggr_kwargs \u001b[38;5;241m=\u001b[39m res[\u001b[38;5;241m0\u001b[39m] \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(res, \u001b[38;5;28mtuple\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m res\n\u001b[1;32m    351\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maggrs) \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 352\u001b[0m     out \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maggregate\u001b[49m\u001b[43m(\u001b[49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43maggr_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    353\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    354\u001b[0m     outs \u001b[38;5;241m=\u001b[39m []\n",
      "Input \u001b[0;32mIn [9]\u001b[0m, in \u001b[0;36mGINLAFConv.aggregate\u001b[0;34m(self, inputs, index)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21maggregate\u001b[39m(\u001b[38;5;28mself\u001b[39m, inputs, index):\n\u001b[1;32m     10\u001b[0m     x \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39msigmoid(inputs)\n\u001b[0;32m---> 11\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaf\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     12\u001b[0m     x \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mview((\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdim \u001b[38;5;241m*\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39munits))\n\u001b[1;32m     13\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmlp(x)\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch/nn/modules/module.py:1110\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1107\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1108\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1109\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1110\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1112\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36mLAFLayer.forward\u001b[0;34m(self, data, index, dim, **kwargs)\u001b[0m\n\u001b[1;32m     74\u001b[0m exps \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39munsqueeze(exps, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[1;32m     75\u001b[0m exps \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mpow(exps, torch\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mweights[\u001b[38;5;241m0\u001b[39m:\u001b[38;5;241m4\u001b[39m]))\n\u001b[0;32m---> 77\u001b[0m scatter \u001b[38;5;241m=\u001b[39m \u001b[43mtorch_scatter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mscatter_add\u001b[49m\u001b[43m(\u001b[49m\u001b[43mexps\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mview\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdim\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     78\u001b[0m scatter \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mclamp(scatter, eps)\n\u001b[1;32m     80\u001b[0m sqrt \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mpow(scatter, torch\u001b[38;5;241m.\u001b[39mrelu(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mweights[\u001b[38;5;241m4\u001b[39m:\u001b[38;5;241m8\u001b[39m]))\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch_scatter/scatter.py:29\u001b[0m, in \u001b[0;36mscatter_add\u001b[0;34m(src, index, dim, out, dim_size)\u001b[0m\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mscatter_add\u001b[39m(src: torch\u001b[38;5;241m.\u001b[39mTensor, index: torch\u001b[38;5;241m.\u001b[39mTensor, dim: \u001b[38;5;28mint\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m     27\u001b[0m                 out: Optional[torch\u001b[38;5;241m.\u001b[39mTensor] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m     28\u001b[0m                 dim_size: Optional[\u001b[38;5;28mint\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m torch\u001b[38;5;241m.\u001b[39mTensor:\n\u001b[0;32m---> 29\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mscatter_sum\u001b[49m\u001b[43m(\u001b[49m\u001b[43msrc\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindex\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdim_size\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/anaconda3/envs/geometric_new/lib/python3.9/site-packages/torch_scatter/scatter.py:19\u001b[0m, in \u001b[0;36mscatter_sum\u001b[0;34m(src, index, dim, out, dim_size)\u001b[0m\n\u001b[1;32m     17\u001b[0m     size[dim] \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m---> 19\u001b[0m     size[dim] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mint\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmax\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m     20\u001b[0m out \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mzeros(size, dtype\u001b[38;5;241m=\u001b[39msrc\u001b[38;5;241m.\u001b[39mdtype, device\u001b[38;5;241m=\u001b[39msrc\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m out\u001b[38;5;241m.\u001b[39mscatter_add_(dim, index, src)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "net = \"LAF\"\n",
    "if net == \"LAF\":\n",
    "    model = LAFNet().to(device)\n",
    "elif net == \"PNA\":\n",
    "    model = PNANet().to(device)\n",
    "elif net == \"GIN\":\n",
    "    GINNet().to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "\n",
    "def train(epoch):\n",
    "    model.train()\n",
    "\n",
    "    if epoch == 51:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.5 * param_group['lr']\n",
    "\n",
    "    loss_all = 0\n",
    "    for data in train_loader:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data.x, data.edge_index, data.batch)\n",
    "        loss = F.nll_loss(output, data.y)\n",
    "        loss.backward()\n",
    "        loss_all += loss.item() * data.num_graphs\n",
    "        optimizer.step()\n",
    "    return loss_all / len(train_dataset)\n",
    "\n",
    "\n",
    "def test(loader):\n",
    "    model.eval()\n",
    "\n",
    "    correct = 0\n",
    "    for data in loader:\n",
    "        data = data.to(device)\n",
    "        output = model(data.x, data.edge_index, data.batch)\n",
    "        pred = output.max(dim=1)[1]\n",
    "        correct += pred.eq(data.y).sum().item()\n",
    "    return correct / len(loader.dataset)\n",
    "\n",
    "\n",
    "for epoch in range(1, 101):\n",
    "    train_loss = train(epoch)\n",
    "    train_acc = test(train_loader)\n",
    "    test_acc = test(test_loader)\n",
    "    print('Epoch: {:03d}, Train Loss: {:.7f}, '\n",
    "          'Train Acc: {:.7f}, Test Acc: {:.7f}'.format(epoch, train_loss,\n",
    "                                                       train_acc, test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
